##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(RColorBrewer)
library(ggpubr)
library(ggsci)

##########################################################################################

option_list <- list(
    make_option(c("--muti_cancer"), type = "character") ,
    make_option(c("--muti_pre"), type = "character") ,
    make_option(c("--gene_list"), type = "character") ,
    make_option(c("--sample_info"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    muti_cancer <- "~/20220915_gastric_multiple/dna_combinePublic/maf/All_GGA.cancer.maf"
    muti_pre <- "~/20220915_gastric_multiple/dna_combinePublic/maf/All_GGA.precancer.maf"
    sample_info <- "~/20220915_gastric_multiple/dna_combinePublic/config/tumor_normal.class.list"
    gene_list <- "~/20220915_gastric_multiple/dna_combinePublic/public_ref/importTantGene.list"
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/images/ITH"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

gene_list <- opt$gene_list
muti_cancer <- opt$muti_cancer
muti_pre <- opt$muti_pre
sample_info <- opt$sample_info
out_path <- opt$out_path

###########################################################################################

dir.create(out_path , recursive = T)
col <- c( "#006699","#DDA520"  )

###########################################################################################

dat_mutiCancer <- fread(muti_cancer  ,sep = "\t",  quote="" ,header = T)
dat_mutiPre <- fread(muti_pre  ,sep = "\t",  quote="" ,header = T)
dat_mutiNodule <- rbind(dat_mutiCancer , dat_mutiPre) 

dat_info <- data.frame(fread(sample_info))
dat_gene <- data.frame(fread(gene_list))

###########################################################################################

col<-brewer.pal(12,"Set3")
col_point <- brewer.pal(9,"Set1")
Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################
dat_mutiNodule <- subset(dat_mutiNodule , t_alt_count > 0)

dat_mutiNodule <- data.frame(Hugo_Symbol = dat_mutiNodule$Hugo_Symbol,
    Chromosome = dat_mutiNodule$Chromosome , Start_Position =  dat_mutiNodule$Start_Position , End_Position = dat_mutiNodule$End_Position ,
    Reference_Allele = dat_mutiNodule$Reference_Allele , Tumor_Seq_Allele2 = dat_mutiNodule$Tumor_Seq_Allele2 , 
    Variant_Classification = dat_mutiNodule$Variant_Classification,
    t_ref_count = dat_mutiNodule$t_ref_count , t_alt_count = dat_mutiNodule$t_alt_count ,
    Tumor_Sample_Barcode = dat_mutiNodule$Tumor_Sample_Barcode)

dat_mutiNodule$Location <- paste( dat_mutiNodule$Chromosome , dat_mutiNodule$Start_Position , 
    dat_mutiNodule$Reference_Allele , dat_mutiNodule$Tumor_Seq_Allele2 , sep=":"  )

###########################################################################################
## 计算每个Normal的Tumor的瘤内异质性
## 共享的突变数目/各自私有的突变数目之和
## Genomic comparison of esophageal squamous cell carcinoma and its precursor lesions by multi-region whole-exome sequencing
## 2017. Nature Communications
###########################################
## Function
computeITH <- function( tmp1 , tmp2 ) {
    shareMut_Num <- length(which(tmp1$Location %in% tmp2$Location))
    t1_private_Num <- nrow(tmp1[!(tmp1$Location %in% tmp2$Location),])
    t2_private_Num <- nrow(tmp2[!(tmp2$Location %in% tmp1$Location),])

    ITH <- 1 - shareMut_Num/(t1_private_Num + t2_private_Num + shareMut_Num)

    res_tmp <- data.frame( shareMut_Num = shareMut_Num , t1_private_Num = t1_private_Num , t2_private_Num = t2_private_Num , ITH = ITH )

    return(res_tmp)
}

###########################################
resulst_hetetogeneity <- c()

for( Sample in unique(dat_info$ID) ){
    tmp <- subset( dat_info , ID == Sample )

    ## 计算ITH
    for( i in 1:(dim(tmp)[1]-1) ){
        for( j in (i+1):dim(tmp)[1] ){
            tumor1 <- tmp[i , "Tumor"]
            tumor2 <- tmp[j , "Tumor"]

            class1 <- tmp[i , "Class"]
            class2 <- tmp[j , "Class"]

            tmp1 <- subset( dat_mutiNodule , Tumor_Sample_Barcode == tumor1 )
            tmp2 <- subset( dat_mutiNodule , Tumor_Sample_Barcode == tumor2 )

            res_tmp <- computeITH( tmp1 , tmp2 )

            res_tmp2 <- cbind(data.frame( Normal = Sample , Tumor_1 = tumor1 , Tumor_2 = tumor2 ,Class_1 = class1 , Class_2 = class2) , res_tmp )

            resulst_hetetogeneity <- rbind( resulst_hetetogeneity , res_tmp2 )
        }
    }
}


###########################################################################################
## 判断肿瘤的类型
## 1、无驱动突变
## 2、已报道的driver突变出现在Trunk
## 3、driver突变均在分支
result_driver <- c()

for( Sample in unique(dat_info$ID) ){

    tumors <- subset( dat_info , ID == Sample )$Tumor

    ## 确定driver突变
    tmp <- subset( dat_mutiNodule , Tumor_Sample_Barcode %in% tumors & Variant_Classification %in% Variant_Type & Hugo_Symbol %in% dat_gene$Gene_Symbol )
    ## 判断突变的数量
    tmp_driver <- tmp %>% 
    group_by(Location) %>%
    summarize( MutTumor = length(Tumor_Sample_Barcode) )

    if( nrow(tmp_driver) == 0){
        class <- "NoDriver"
    }else if( nrow(tmp_driver) > 0 ){
        share_driver <- length(which(tmp_driver$MutTumor == length(tumors)))
        if(share_driver > 0 ){
            class <- "ShareDriver"
        }else{
            class <- "PrivateDriver"
        }
    }

    tmp_res <- data.frame( Normal = Sample , DriverClass = class )

    result_driver <- rbind( result_driver , tmp_res )

}

## 合并异质性和驱动突变
result <- merge( resulst_hetetogeneity , result_driver , by = "Normal" )

## 描述Tumor的Class
pre_class <- c("IM")
cancer_class <- c("DGC" , "IGC")

result$TypeClass <- ""
result$TypeClass <- ifelse( 
    (result$Class_1 %in% pre_class & result$Class_2 %in% cancer_class ) | (result$Class_2 %in% pre_class & result$Class_1 %in% cancer_class) , 
    "Pre_Inv", result$TypeClass)
result$TypeClass <- ifelse(result$Class_1 %in% pre_class & result$Class_2 %in% pre_class , "Pre_Pre", result$TypeClass)
result$TypeClass <- ifelse(result$Class_1 %in% cancer_class & result$Class_2 %in% cancer_class , "Inv_Inv", result$TypeClass)
result$TypeClass <- factor( result$TypeClass , levels = c("Pre_Pre" , "Pre_Inv" , "Inv_Inv" ) , order = T )

## 区分IGC和DGC，还有IM
result$TypeClass2 <- ""
result$TypeClass2 <- ifelse( 
    (result$Class_1 %in% "IM" & result$Class_2 %in% "DGC" ) | (result$Class_2 %in% "IM" & result$Class_1 %in% "DGC") , 
    "IM_DGC", result$TypeClass2)
result$TypeClass2 <- ifelse( 
    (result$Class_1 %in% "IM" & result$Class_2 %in% "IGC" ) | (result$Class_2 %in% "IM" & result$Class_1 %in% "IGC") , 
    "IM_IGC", result$TypeClass2)
result$TypeClass2 <- ifelse(result$Class_1 %in% "IM" & result$Class_2 %in% "IM" , "IM_IM", result$TypeClass2)
result$TypeClass2 <- factor( result$TypeClass2 , levels = c("IM_IM" , "IM_DGC" , "IM_IGC" ) , order = T )

col_point <- col_point[3:1]
names(col_point) <- c("Pre_Pre" , "Pre_Inv" , "Inv_Inv" )

###########################################################################################
## 合并无驱动突变 + driver突变出现在分支
## 1、已报道的driver突变出现在Trunk
## 2、无Trunk突变
## 其ITH有无差异
my_comparisons <- list( c("NoTrunkDriver", "TrunkDriver"))
result$DriverClass_combine <- ifelse( result$DriverClass=="ShareDriver" , "TrunkDriver" , "NoTrunkDriver")
result$DriverClass_combine <- factor( result$DriverClass_combine , levels = c( "TrunkDriver" , "NoTrunkDriver") , order = T )

result$TypeClass_combine <- ifelse( result$TypeClass %in% c("Pre_Inv" , "Pre_Pre") , "IM_IM|GC" , "GC_GC" )
result$TypeClass_combine <- factor( result$TypeClass_combine , levels = c( "IM_IM|GC" , "GC_GC") , order = T )
dat_tmp <- merge( result , unique(dat_info[,c("ID" , "Type")]) , by.x = "Normal" , by.y = "ID")
dat_tmp2 <- dat_tmp
dat_tmp2$Type <- "All"
dat_plot <- rbind( dat_tmp , dat_tmp2 )
dat_plot$Type <- factor( dat_plot$Type , levels = c("All" , "IM + IGC" , "IM + DGC") , order = T )

plot <- ggplot(data=result,mapping = aes(x=DriverClass_combine,y=ITH))+
  geom_boxplot(lwd=1.5,aes(color=DriverClass_combine)) +
  geom_jitter(position=position_jitter(0.2),aes(color=DriverClass_combine)) +
  scale_color_npg() +
  #facet_grid(.~Type)+
  xlab(NULL) +
  ylab('ITH')+
  theme_bw() +
  ylim(-0.1,1.3) +
  stat_compare_means(aes(group = DriverClass_combine, label = sprintf("p = %2.1e", as.numeric(..p.format..))) , method = "wilcox.test" , size = 6 , label.y = 1.2 ) +
  #stat_compare_means(label.y = 1.1) +
  theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='right',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        legend.text = element_text(size = 8,color="black",face='bold'),
        axis.text.x = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 15,color="black",face='bold'),
        axis.title.x = element_text(size = 18,color="black",face='bold'),
        axis.title.y = element_text(size = 18,color="black",face='bold'),
        axis.line = element_line(size = 0.5)) 

out_name <- paste0(out_path , "/ITH.combine.TreeClass.pdf")
ggsave(file=out_name,plot=plot,width=5,height=6)


## 统计Pre_Pre、Pre_Inv、Inv_Inv
## 在三类之间的分布
tmp_res <- matrix(table(paste0(result$DriverClass_combine,":" , result$TypeClass_combine)) , ncol=2)
colnames(tmp_res) <- c("NoTrunkDriver" , "TrunkDriver")
rownames(tmp_res) <- c("Pre_Pre|Inv" , "Inv_Inv")
tmp_res <- data.frame(t(tmp_res))

out_name <- paste0(out_path , "/ITH_typeClass.combine.tsv")
write.table( tmp_res , out_name , row.names = T , sep = "\t" , quote = F )


###########################################################################################
## 输出计算的瘤内异质性
out_name <- paste0(out_path , "/ITH.tsv")
write.table( result , out_name , row.names = F , sep = "\t" , quote = F )

## 判断每个人是否共享驱动突变
result_normal <- data.frame(Normal = result$Normal , DriverClass = result$DriverClass , DriverClass_combine = result$DriverClass_combine)
result_normal <- unique(result_normal)
out_name <- paste0(out_path , "/DriverClass.tsv")
write.table( result_normal , out_name , row.names = F , sep = "\t" , quote = F )


###########################################################################################
## 计算TrunkDriver和NoTrunkDriver的两组亚型的瘤内异质性
my_comparisons <- list( c("NoTrunkDriver", "TrunkDriver"))

out_name <- paste0(out_path , "/ITH.combine.TreeClass.ThreeSubClass.pdf")

plot <- ggplot(data=result,mapping = aes(x=TypeClass_combine,y=ITH))+
    geom_boxplot(lwd=1.5,aes(color=DriverClass_combine)) +
    geom_point(position = position_jitterdodge(),aes(color=DriverClass_combine)) +
    scale_color_npg() +
    xlab(NULL) +
    ylab('ITH')+
    theme_bw() +
    ylim(-0.1,1.3) +
    stat_compare_means(aes(group = DriverClass_combine, label = sprintf("p = %2.1e", as.numeric(..p.format..))) , method = "wilcox.test" , size = 4 , label.y = 1.2   ) +
    stat_compare_means(aes(group = TypeClass_combine, label = sprintf("p = %2.1e", as.numeric(..p.format..))) , method = "wilcox.test" , size = 5 , label.y = 1.3 , label.x=1.2 ) +
    #stat_compare_means(label.y = 1.1) +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='right',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        legend.text = element_text(size = 8,color="black",face='bold'),
        axis.text.x = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 15,color="black",face='bold'),
        axis.title.x = element_text(size = 18,color="black",face='bold'),
        axis.title.y = element_text(size = 18,color="black",face='bold'),
        axis.line = element_line(size = 0.5)) 
ggsave(file=out_name,plot=plot,width=5,height=6)


my_comparisons <- list( c("Pre_Pre|Inv" , "Inv_Inv"))

out_name <- paste0(out_path , "/ITH.combine.TreeClass.TrunkSubClass.driver.pdf")

plot <- ggplot(data=result,mapping = aes(x=DriverClass_combine,y=ITH))+
    geom_boxplot(lwd=1.5,aes(color=TypeClass_combine)) +
    geom_point(position = position_jitterdodge(),aes(color=TypeClass_combine)) +
    scale_color_npg() +
    xlab(NULL) +
    ylab('ITH')+
    theme_bw() +
    ylim(-0.1,1.3) +
    stat_compare_means(aes(group = TypeClass_combine, label = sprintf("p = %2.1e", as.numeric(..p.format..))) , method = "wilcox.test" , size = 4 , label.y = 1.2   ) +
    stat_compare_means(aes(group = DriverClass_combine, label = sprintf("p = %2.1e", as.numeric(..p.format..))) , method = "wilcox.test" , size = 5 , label.y = 1.3 , label.x=1.2  ) +
    #stat_compare_means(label.y = 1.1) +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='right',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        legend.text = element_text(size = 8,color="black",face='bold'),
        axis.text.x = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 15,color="black",face='bold'),
        axis.title.x = element_text(size = 18,color="black",face='bold'),
        axis.title.y = element_text(size = 18,color="black",face='bold'),
        axis.line = element_line(size = 0.5)) 
ggsave(file=out_name,plot=plot,width=5,height=6)


###########################################################################################
## IM和IGC
## IM个DGC

my_comparisons <- list( c("IM_IM", "IM_IGC") , c( "IM_IM", "IM_DGC") , c("IM_IGC" , "IM_DGC"))

out_name <- paste0(out_path , "/ITH.combine.laruenClass.pdf")

plot <- ggplot(data=result[!is.na(result$TypeClass2),],mapping = aes(x=TypeClass2,y=ITH))+
    geom_boxplot(lwd=1.5,aes(color=TypeClass2)) +
    geom_point(position = position_jitterdodge(),aes(color=TypeClass2)) +
    scale_color_npg() +
    #scale_color_manual(values=col) +
    xlab(NULL) +
    ylab('ITH')+
    theme_bw() +
    ylim(-0.1,1.3) +
    stat_compare_means(comparisons = my_comparisons) +
      #stat_compare_means(label.y = 1.1) +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='right',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        legend.text = element_text(size = 8,color="black",face='bold'),
        axis.text.x = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 15,color="black",face='bold'),
        axis.title.x = element_text(size = 18,color="black",face='bold'),
        axis.title.y = element_text(size = 18,color="black",face='bold'),
        axis.line = element_line(size = 0.5)) 
ggsave(file=out_name,plot=plot,width=5,height=6)
